import os
import cv2
import pynvml
import warnings
#! 设置显卡
pynvml.nvmlInit() #! 初始化
deviceCount = pynvml.nvmlDeviceGetCount() #! 统计显卡数量
GPUs = []
for busid in range(deviceCount):
    handle = pynvml.nvmlDeviceGetHandleByIndex(busid)
    device = pynvml.nvmlDeviceGetMemoryInfo(handle)
    memory = device.total // 1048576 / 1024 #! 单位: G
    if memory > 8: GPUs.append(str(busid))
pynvml.nvmlShutdown() #! 关闭
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(GPUs) #! 设置显卡设备
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" #! 设置警告ERROR级别
#! 导入官方包
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
#! 导入自定义模块
from model import SSD300, SSD512
from data import PascalVOCDatasetV1, PascalVOCDatasetV2
#! 忽略警告类报错
warnings.filterwarnings('ignore')

#! 数据集设置
data_folder = 'dataset/VOCdevkit'
voc_labels = ('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
              'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor')
label_map = {k: v + 1 for v, k in enumerate(voc_labels)}
label_map['background'] = 0
n_classes = len(label_map)
rev_label_map = {v:k for k,v in label_map.items()}
batch_size = 8
output_shape = (300,300)
PascalVOCDataset = PascalVOCDatasetV1
train_loader = PascalVOCDataset(batch_size,data_folder, 'TRAIN', output_shape=output_shape, shuffle=True)
valid_loader = PascalVOCDataset(batch_size,data_folder, 'TEST', output_shape=output_shape)

#! 模型设置
model_path = None
SSD = SSD300
SSDModel = SSD(n_classes=n_classes, model_path=model_path) if model_path is not None else SSD(n_classes=n_classes)

#! 模型评估
SSDModel.evaluate(data_loader=valid_loader)

#! 预测
image_paths, image_bboxes, image_labels, image_scores = SSDModel.predict(min_score=0.05, max_overlap=0.6, top_k=100, data_loader=valid_loader)

#! 绘图展示
for image_path, bboxes, labels, scores in zip(image_paths, image_bboxes, image_labels, image_scores):
    
    image = cv2.imread(image_path)
    h, w, c = image.shape
    bboxes = (bboxes * np.array([w,h,w,h])).astype(np.int64)

    for box, label, score in zip(bboxes, labels, scores):
        if score > 0.45:
            x1, y1, x2, y2 = box
            image = cv2.rectangle(image,(x1,y1),(x2,y2),(255,0,0))
            print(rev_label_map[int(label)], score)

    print("-=-"*30)
    cv2.imshow('image',image)
    cv2.waitKey()
    cv2.destroyAllWindows()

    break